# -*- coding: utf8 -*-

import os, datetime, string, sys
import time, threading
import logging, logging.handlers
import socket, SocketServer
from yarascanner import *
from icapserver import *

def dump(obj):
    for attr in dir(obj):
        print "obj.%s = %s" % (attr, getattr(obj, attr))

yara = YaraScanner()

class ThreadingSimpleServer(SocketServer.ThreadingMixIn, ICAPServer):
    pass

class YaraICAPHandler(BaseICAPRequestHandler):

    def yara_OPTIONS(self):
        self.set_icap_response(200)
        self.set_icap_header('Methods', 'REQMOD')
        self.set_icap_header('Service', 'ICAP Server' + ' ' + self._server_version)
        self.set_icap_header('Options-TTL', '3600')
        self.set_icap_header('Preview', '0')
        self.send_headers(False)

    def yara_REQMOD(self):
        #self.no_adaptation_required() 
        if '204' in self.allow or self.preview != None:
			# We MUST read everything the client sent us
            if self.has_body:
                while True:
                    if self.read_chunk() == '':
                        break
                    
            content = self.enc_req[1]
            ip = str(self.enc_req_headers['host'][0])
            port = str(self.enc_req[1])[:str(self.enc_req[1]).find(':')]
            if (ip, port) in yara_rule.yara_rule:
                action = [1]
                if self.headers.has_key('x-client-ip'):
                    yara.Scan(content, self.enc_req, self.enc_req_headers, self.enc_res_headers, self.headers['x-client-ip'], action)
                else:
                    yara.Scan(content, self.enc_req, self.enc_req_headers, self.enc_res_headers, ["anon-client-ip"], action)

                #add drop packet (0 drop 1 permit)
                if action[0] == 0:
                    self.set_icap_response(400)
                    self.send_headers()
                    return
                else:
                    content = ''               
                    enc_req = self.enc_req[0] + self.enc_req[1] + self.enc_req[2]
                    content += enc_req
                    
                    for i, val in enumerate(self.enc_req_headers.values()):
                        for j, valval in enumerate(val):
                            content += valval

                    if self.headers.has_key('x-client-ip'):
                        yara.Scan(content, self.enc_req, self.enc_req_headers, self.enc_res_headers, self.headers['x-client-ip'], action)
                    else:
                        yara.Scan(content, self.enc_req, self.enc_req_headers, self.enc_res_headers, ["anon-client-ip"], action)
                    if action[0] == 0:
                        self.set_icap_response(400)
                        self.send_headers()
                        return
            self.set_icap_response(204)
            self.send_headers()

        else:
			# We have to copy everything,
			# but it's sure there's no preview
            self.set_icap_response(200)

            self.set_enc_status(' '.join(self.enc_res_status))
            for h in self.enc_res_headers:
                for v in self.enc_res_headers[h]:
                    self.set_enc_header(h, v)

            if not self.has_body:
                self.send_headers(False)
                return
            
            content = self.enc_req[1]

            ip = str(self.enc_req_headers['host'][0])
            port = str(self.enc_req[1])[:str(self.enc_req[1]).find(':')]
            if (ip, port) in yara_rule.yara_rule:
                action = [1]
                if self.headers.has_key('x-client-ip'):
                    yara.Scan(content, self.enc_req, self.enc_req_headers, self.enc_res_headers, self.headers['x-client-ip'], action)
                else:
                    yara.Scan(content, self.enc_req, self.enc_req_headers, self.enc_res_headers, ["anon-client-ip"], action)

                #add drop packet (0 drop 1 permit)
                if action[0] == 0:
                    print(content)
                    self.set_icap_response(400)
                    self.send_headers()
                    return
                else:
                    content = ''               
                    enc_req = self.enc_req[0] + self.enc_req[1] + self.enc_req[2]
                    content += enc_req
                    
                    for i, val in enumerate(self.enc_req_headers.values()):
                        for j, valval in enumerate(val):
                            content += valval

                    if self.headers.has_key('x-client-ip'):
                        yara.Scan(content, self.enc_req, self.enc_req_headers, self.enc_res_headers, self.headers['x-client-ip'], action)
                    else:
                        yara.Scan(content, self.enc_req, self.enc_req_headers, self.enc_res_headers, ["anon-client-ip"], action)
                    if action[0] == 0:
                        self.set_icap_response(400)
                        self.send_headers()
                        return

            self.send_headers(True)
            while True:
                chunk = self.read_chunk()
                self.send_chunk(chunk)
                if chunk == '':
                    break

    def yara_RESPMOD(self):
        self.set_icap_response(200)
        self.set_enc_status(' '.join(self.enc_res_status))
        for h in self.enc_res_headers:
            for v in self.enc_res_headers[h]:
                self.set_enc_header(h, v)

        if not self.has_body:
            self.send_headers(False)
            return

        # self.send_headers(True)
        content = ''
        while True:
            chunk = self.read_chunk()
            # self.send_chunk(chunk)
            content += chunk
            if chunk == '':
                break
        """   
        ip = str(self.enc_req_headers['host'][0])
        port = str(self.enc_req[1])[:str(self.enc_req[1]).find(':')]
        if (ip, port) in yara_rule.yara_rule:
            action = [1]
            if self.headers.has_key('x-client-ip'):
                yara.Scan(content, self.enc_req, self.enc_req_headers, self.enc_res_headers, self.headers['x-client-ip'], action)
            else:
                yara.Scan(content, self.enc_req, self.enc_req_headers, self.enc_res_headers, ["anon-client-ip"], action)
            
            #add drop packet (0 drop 1 permit)
            if action[0] == 0:
                self.set_icap_response(400)
                self.send_headers(False)
                return
        """              
        self.send_headers(True)
        self.send_chunk(content)

try:
    #print 'Use Control-C to exit'
    server = ThreadingSimpleServer(('127.0.0.1', 1344), YaraICAPHandler)
    server.serve_forever()
except KeyboardInterrupt:
    server.shutdown()
    server.server_close()
    #print "Finished"
